# coding: UTF-8
import os
import re
import torch
import jieba

import pandas as pd
import numpy as np
import pkg_resources
import pickle as pkl
from tqdm import tqdm
import time
from datetime import timedelta
from gensim.models import word2vec


MAX_VOCAB_SIZE = 10000  # 词表长度限制
UNK, PAD, NUM = '<UNK>', '<PAD>', '<NUM>'  # 未知字，padding符号
stopwords_path = 'source/rule/stopwords.txt'
model_name = 'w2v_with_cut_d100_20200727_112207'

def read_from_txt(path, ID2LABEL):
    texts = []
    text_labels = []
    with open(path, 'r', encoding='UTF-8') as f:
        # Text_labels after text and a space is in between
        for line in tqdm(f):
            lin = line.strip()
            try:
                l, t, _ = lin.split('\t')
            except:
                print(lin)
                input()
            if int(l) == 2 or int(l) == 3:
                continue
            texts.append(t)
            text_labels.append(ID2LABEL[int(l)])
    return texts, text_labels

def load_w2v_model():
    starttime = time.time()
    w2v_model_path = pkg_resources.resource_filename(__name__, 'source/word2vector/'+model_name)
    w2v_model = word2vec.Word2Vec.load(w2v_model_path)
    endtime = time.time()
    print('加载 w2v 模型耗时：%.2f s' %(endtime-starttime,))
    return w2v_model
w2v_model = load_w2v_model()

def build_vocab(texts, tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=2):
    vocab_dic = {}
    for text in tqdm(texts): 
        for word in tokenizer(text):
            vocab_dic[word] = vocab_dic.get(word, 0) + 1
    vocab_list = (sorted([_ for _ in vocab_dic.items() if _[1] >= min_freq], 
                  key=lambda x: x[1], reverse=True)[:max_size])
    vocab_dic = {word_count[0]: idx for idx, word_count in enumerate(vocab_list)}
    vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})
    print(vocab_dic)
    return vocab_dic


def convert2tensor(texts, tokenizer, vocab, device, pad_size=32):
    # 把所有句子补长到pad_size=50, 超过则截短到50
    contents = []
    lengths = []
    
    '''
    while True:
        word = input()
        try:
            res = w2v_model.wv.most_similar(word,topn=10)
        except:
            res = 'Key %s not present' %(word)
        print(res)
    '''
    for text in tqdm(texts):
        tokens = tokenizer(text)
        seq_len = len(tokens)
        if len(tokens) < pad_size:
            tokens.extend([PAD] * (pad_size - len(tokens)))
        else:
            tokens = tokens[:pad_size]
            seq_len = pad_size
        words_line = []
        for token in tokens:
            try:
                _ = vocab[token]
                words_line.append(list(w2v_model.wv[token]))
            except:
                words_line.append(list(np.zeros((100))))
        contents.append(words_line)
        lengths.append(seq_len)
    contents = torch.FloatTensor(contents).to(device)
    lengths = torch.LongTensor(lengths).to(device)
    return contents, lengths


def get_time_dif(start_time):
    end_time = time.time()
    time_dif = end_time - start_time
    return timedelta(seconds=int(round(time_dif)))

def tokenizer(text):
    '''分词+去停用词'''
    stopwords = get_stopwords_list()
    sentence_depart = seg_depart(text)
    sentence_depart = remove_stopwords(sentence_depart, stopwords)
    return sentence_depart

def get_stopwords_list():
    stopwords = [line.strip() for line in open(stopwords_path, encoding='UTF-8').readlines()]
    return stopwords

def seg_depart(sentence):
    sentence_depart = jieba.lcut(sentence.strip(), cut_all=False, HMM=True)
    return sentence_depart

def remove_digits(input_str):
    punc = u'0123456789.'
    output_str = re.sub(r'[{}]+'.format(punc), '', input_str)
    return output_str

def remove_stopwords(sentence_list, stopwords_list):
    out_list = []
    for word in sentence_list:
        if word not in stopwords_list:
            if not remove_digits(word):
                continue
            if word != '\t':
                out_list.append(word)
    return out_list


if __name__ == "__main__":
    sentence = '1、判令被告赵军霞偿还原告借款本息及应收费用共计4278.6元（计算至2017年1月10日，实际还款额以合同约定的计费方式计算至最终还款日）'
    
    input()
    '''
    tokenizer = lambda x: [y for y in x]
    print(build_vocab('news.txt', tokenizer, MAX_VOCAB_SIZE, 2))
    '''
    # 下面的目录、文件名按需更改。
    train_dir = "./THUCNews/data/train.txt"
    vocab_dir = "./THUCNews/data/vocab.pkl"
    pretrain_dir = "./THUCNews/data/sgns.sogou.char"
    emb_dim = 300
    filename_trimmed_dir = "./THUCNews/data/embedding_SougouNews"
    if os.path.exists(vocab_dir):
        word_to_id = pkl.load(open(vocab_dir, 'rb'))
    else:
        # tokenizer = lambda x: x.split(' ')  # 以词为单位构建词表(数据集中词之间以空格隔开)
        tokenizer = lambda x: [y for y in x]  # 以字为单位构建词表
        word_to_id = build_vocab(train_dir, tokenizer=tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=1)
        pkl.dump(word_to_id, open(vocab_dir, 'wb'))

    embeddings = np.random.rand(len(word_to_id), emb_dim)
    f = open(pretrain_dir, "r", encoding='UTF-8')
    for i, line in enumerate(f.readlines()):
        # if i == 0:  # 若第一行是标题，则跳过
        #     continue
        lin = line.strip().split(" ")
        if lin[0] in word_to_id:
            idx = word_to_id[lin[0]]
            emb = [float(x) for x in lin[1:301]]
            embeddings[idx] = np.asarray(emb, dtype='float32')
    f.close()
    np.savez_compressed(filename_trimmed_dir, embeddings=embeddings)
    
